import xml.etree.ElementTree as ET
import os
import shutil
import cv2

import imgaug as ia
from imgaug import augmenters as iaa
from utils.basic import cv_imread, cv_imwrite


def load_from_xml_file(xml_file_path: str):
    '''
     函数功能:加载xml文件，将xml里的信息保存好
    :return:
    '''
    tree = ET.parse(xml_file_path)
    root = tree.getroot()
    boundingboxes = []
    for obj in root.iter('object'):
        class_name = obj.find('./name').text
        difficult = bool(int(obj.find('difficult').text))
        xml_box = obj.find('bndbox')
        _xmin = int(xml_box.find('xmin').text)
        _xmax = int(xml_box.find('xmax').text)
        _ymin = int(xml_box.find('ymin').text)
        _ymax = int(xml_box.find('ymax').text)
        boundingboxes.append([_xmin, _ymin, _xmax, _ymax])
    return boundingboxes


def change_from_xml_file(xml_file_path: str, after_boundingboxes: list, shape=None):
    '''
     函数功能:加载xml文件，将xml里的信息保存好
    :return:
    '''
    tree = ET.parse(xml_file_path)
    root = tree.getroot()
    size = root.find('size')
    height = size.find('height')
    width = size.find('width')
    if shape != None:
        height.text = str(shape[0])
        width.text = str(shape[1])

    for idx, obj in enumerate(root.iter('object')):
        xml_box = obj.find('bndbox')
        xml_box.find('xmin').text = str(after_boundingboxes[idx][0])
        xml_box.find('xmax').text = str(after_boundingboxes[idx][2])
        xml_box.find('ymin').text = str(after_boundingboxes[idx][1])
        xml_box.find('ymax').text = str(after_boundingboxes[idx][3])
    tree.write(xml_file_path)


def scale(path: str, save_path: str, scale_size, interpolation='cubic'):
    '''
    将图像缩放到固定大小
    :param path:要被缩放的路径
    :param save_path:缩放后保存的路径
    :param scale_size:字符串”keep”,此时保持图像原大小不坐缩放。如果是一个整数n,则缩放成(n, n)。
    如果是一个float v,则每张图片会被缩放成(H*v, W*v),此时每张图像大小仍然不一样。
    如果是一个tuple类型(a, b), 如果a、b中至少有一个小数,则从[a,b]中挑选一个数作为缩放比例。
    如果a、b都是整数,则从[a,b]中挑选一个整数作为缩放后的大小。如果是1个list,则list中的数要么全是整数,
    要么全是小数(不能混用)。如果是一个dict类型,则该dict必须有两个key: height和width。
    每个key的值仍然可以按照上面的方法选取。此外, key的值还可以是”keep-aspect-ratio”, 表示按照比例缩放。
    :param interpolation:缩放方法。如果是All, 则会随机从下面挑选一个 : nearest、linear、area、cubic,
    注意每张图片可能不一样。如果是int,则应该是下面的一种：cv2.INTER_NEAREST, cv2.INTER_LINEAR,
    cv2.INTER_AREA,cv2.INTER_CUBIC。如果是string,则该种方法会被一直使用,必须是下面的一种
    ： nearest, linear, area, cubic。如果是int list或者string list, 则每张图片会从中随机选取一个。
    '''
    for idx, image_name in enumerate(os.listdir(path)):
        if image_name.split('.')[-1] == 'png' or image_name.split('.')[-1] == 'jpg':
            image_path = os.path.join(path, image_name)
            xml_path = os.path.join(path, image_name.split('.')[0] + '.xml')
            img = cv_imread(image_path)
            boundingboxes = load_from_xml_file(xml_path)
            print(img.shape)
            seq = iaa.Sequential([
                iaa.Scale(size=scale_size,
                          interpolation=interpolation,
                          name=None,
                          deterministic=False,
                          random_state=None,
                          ),  # vertically flip 20% of all image
            ])
            seq_det = seq.to_deterministic()  # 保持坐标和图像同步改变，而不是随机
            before_boundingBoxs = []
            for boundingbox in boundingboxes:
                # bndbox 坐标增强
                before_boundingBoxs.append(ia.BoundingBox(x1=boundingbox[0], y1=boundingbox[1]
                                                          , x2=boundingbox[2], y2=boundingbox[3]))

            bbs = ia.BoundingBoxesOnImage(before_boundingBoxs, shape=img.shape)
            bbs_aug = seq_det.augment_bounding_boxes([bbs])[0]
            image_aug = seq_det.augment_images([img])[0]
            after_boundingboxes = []

            for i in range(len(bbs.bounding_boxes)):
                after = bbs_aug.bounding_boxes[i]
                after_boundingboxes.append([int(after.x1), int(after.y1), int(after.x2), int(after.y2)])

            for i in range(len(bbs.bounding_boxes)):
                before = bbs.bounding_boxes[i]
                after = bbs_aug.bounding_boxes[i]
                print("BB %d: (%.4f, %.4f, %.4f, %.4f) -> (%.4f, %.4f, %.4f, %.4f)" % (
                    i,
                    before.x1, before.y1, before.x2, before.y2,
                    int(after.x1), int(after.y1), int(after.x2), int(after.y2))
                      )
            cv_imwrite(os.path.join(save_path, image_name), image_aug)
            shutil.copy(xml_path, os.path.join(save_path, image_name.split('.')[0] + '.xml'))
            change_from_xml_file(os.path.join(save_path, image_name.split('.')[0] + '.xml'), after_boundingboxes,
                                 image_aug.shape)


def fliplr(path: str, save_path: str, percent=1.0):
    '''
    将图像水平镜面翻转
    :param path:要被缩放的路径
    :param save_path:缩放后保存的路径
    :param percent:int或者float,每张图片呗翻转的概率
    '''
    for idx, image_name in enumerate(os.listdir(path)):
        if image_name.split('.')[-1] == 'png' or image_name.split('.')[-1] == 'jpg':
            image_path = os.path.join(path, image_name)
            xml_path = os.path.join(path, image_name.split('.')[0] + '.xml')
            img = cv_imread(image_path)
            boundingboxes = load_from_xml_file(xml_path)
            print(img.shape)
            seq = iaa.Sequential([
                iaa.Fliplr(p=percent, name=None, deterministic=False, random_state=None)
            ])
            seq_det = seq.to_deterministic()  # 保持坐标和图像同步改变，而不是随机
            before_boundingBoxs = []
            for boundingbox in boundingboxes:
                # bndbox 坐标增强
                before_boundingBoxs.append(ia.BoundingBox(x1=boundingbox[0], y1=boundingbox[1]
                                                          , x2=boundingbox[2], y2=boundingbox[3]))

            bbs = ia.BoundingBoxesOnImage(before_boundingBoxs, shape=img.shape)
            bbs_aug = seq_det.augment_bounding_boxes([bbs])[0]
            image_aug = seq_det.augment_images([img])[0]
            after_boundingboxes = []

            for i in range(len(bbs.bounding_boxes)):
                after = bbs_aug.bounding_boxes[i]
                after_boundingboxes.append([int(after.x1), int(after.y1), int(after.x2), int(after.y2)])

            for i in range(len(bbs.bounding_boxes)):
                before = bbs.bounding_boxes[i]
                after = bbs_aug.bounding_boxes[i]
                print("BB %d: (%.4f, %.4f, %.4f, %.4f) -> (%.4f, %.4f, %.4f, %.4f)" % (
                    i,
                    before.x1, before.y1, before.x2, before.y2,
                    int(after.x1), int(after.y1), int(after.x2), int(after.y2))
                      )
            cv_imwrite(os.path.join(save_path, image_name), image_aug)
            shutil.copy(xml_path, os.path.join(save_path, image_name.split('.')[0] + '.xml'))
            change_from_xml_file(os.path.join(save_path, image_name.split('.')[0] + '.xml'), after_boundingboxes,
                                 image_aug.shape)


def flipud(path: str, save_path: str, percent=1.0):
    '''
    将图像Flipud
    :param path:要被缩放的路径
    :param save_path:缩放后保存的路径
    :param percent:int或者float,每张图片呗翻转的概率
    '''
    for idx, image_name in enumerate(os.listdir(path)):
        if image_name.split('.')[-1] == 'png' or image_name.split('.')[-1] == 'jpg':
            image_path = os.path.join(path, image_name)
            xml_path = os.path.join(path, image_name.split('.')[0] + '.xml')
            img = cv_imread(image_path)
            boundingboxes = load_from_xml_file(xml_path)
            print(img.shape)
            seq = iaa.Sequential([
                iaa.Flipud(p=percent, name=None, deterministic=False, random_state=None)
            ])
            seq_det = seq.to_deterministic()  # 保持坐标和图像同步改变，而不是随机
            before_boundingBoxs = []
            for boundingbox in boundingboxes:
                # bndbox 坐标增强
                before_boundingBoxs.append(ia.BoundingBox(x1=boundingbox[0], y1=boundingbox[1]
                                                          , x2=boundingbox[2], y2=boundingbox[3]))

            bbs = ia.BoundingBoxesOnImage(before_boundingBoxs, shape=img.shape)
            bbs_aug = seq_det.augment_bounding_boxes([bbs])[0]
            image_aug = seq_det.augment_images([img])[0]
            after_boundingboxes = []

            for i in range(len(bbs.bounding_boxes)):
                after = bbs_aug.bounding_boxes[i]
                after_boundingboxes.append([int(after.x1), int(after.y1), int(after.x2), int(after.y2)])

            for i in range(len(bbs.bounding_boxes)):
                before = bbs.bounding_boxes[i]
                after = bbs_aug.bounding_boxes[i]
                print("BB %d: (%.4f, %.4f, %.4f, %.4f) -> (%.4f, %.4f, %.4f, %.4f)" % (
                    i,
                    before.x1, before.y1, before.x2, before.y2,
                    int(after.x1), int(after.y1), int(after.x2), int(after.y2))
                      )
            cv_imwrite(os.path.join(save_path, image_name), image_aug)
            shutil.copy(xml_path, os.path.join(save_path, image_name.split('.')[0] + '.xml'))
            change_from_xml_file(os.path.join(save_path, image_name.split('.')[0] + '.xml'), after_boundingboxes,
                                 image_aug.shape)


def gaussianBlur(path: str, save_path: str, sigma=3):
    '''
    将图像Flipud
    :param path:要被缩放的路径
    :param save_path:缩放后保存的路径
    :param sigma: 高斯变换的标准差。可为float, float tuple。常见的有0,不扰动。3,强扰动。
    '''
    for idx, image_name in enumerate(os.listdir(path)):
        if image_name.split('.')[-1] == 'png' or image_name.split('.')[-1] == 'jpg':
            image_path = os.path.join(path, image_name)
            xml_path = os.path.join(path, image_name.split('.')[0] + '.xml')
            img = cv_imread(image_path)
            boundingboxes = load_from_xml_file(xml_path)
            print(img.shape)
            seq = iaa.Sequential([
                iaa.GaussianBlur(sigma=sigma, name=None, deterministic=False, random_state=None)
            ])
            seq_det = seq.to_deterministic()  # 保持坐标和图像同步改变，而不是随机
            before_boundingBoxs = []
            for boundingbox in boundingboxes:
                # bndbox 坐标增强
                before_boundingBoxs.append(ia.BoundingBox(x1=boundingbox[0], y1=boundingbox[1]
                                                          , x2=boundingbox[2], y2=boundingbox[3]))

            bbs = ia.BoundingBoxesOnImage(before_boundingBoxs, shape=img.shape)
            bbs_aug = seq_det.augment_bounding_boxes([bbs])[0]
            image_aug = seq_det.augment_images([img])[0]
            after_boundingboxes = []

            for i in range(len(bbs.bounding_boxes)):
                after = bbs_aug.bounding_boxes[i]
                after_boundingboxes.append([int(after.x1), int(after.y1), int(after.x2), int(after.y2)])

            for i in range(len(bbs.bounding_boxes)):
                before = bbs.bounding_boxes[i]
                after = bbs_aug.bounding_boxes[i]
                print("BB %d: (%.4f, %.4f, %.4f, %.4f) -> (%.4f, %.4f, %.4f, %.4f)" % (
                    i,
                    before.x1, before.y1, before.x2, before.y2,
                    int(after.x1), int(after.y1), int(after.x2), int(after.y2))
                      )
            cv_imwrite(os.path.join(save_path, image_name), image_aug)
            shutil.copy(xml_path, os.path.join(save_path, image_name.split('.')[0] + '.xml'))
            change_from_xml_file(os.path.join(save_path, image_name.split('.')[0] + '.xml'), after_boundingboxes,
                                 image_aug.shape)


def sharpen(path: str, save_path: str, alpha=0, lightness=1):
    '''
    将图像锐化
    :param path:要被缩放的路径
    :param save_path:缩放后保存的路径
    :param alpha:number or tuple of number or list of number or imgaug.parameters.StochasticParameter, optional
        Visibility of the sharpened image. At 0, only the original image is
        visible, at 1.0 only its sharpened version is visible.

            * If an int or float, exactly that value will be used.
            * If a tuple ``(a, b)``, a random value from the range ``a <= x <= b`` will
              be sampled per image.
            * If a list, then a random value will be sampled from that list
              per image.
            * If a StochasticParameter, a value will be sampled from the
              parameter per image.
    :param lightness:number or tuple of number or list of number or imgaug.parameters.StochasticParameter, optional
        Parameter that controls the lightness/brightness of the sharped image.
        Sane values are somewhere in the range ``(0.5, 2)``.
        The value 0 results in an edge map. Values higher than 1 create bright
        images. Default value is 1.

            * If an int or float, exactly that value will be used.
            * If a tuple ``(a, b)``, a random value from the range ``a <= x <= b`` will
              be sampled per image.
            * If a list, then a random value will be sampled from that list
              per image.
            * If a StochasticParameter, a value will be sampled from the
              parameter per image.
    '''
    for idx, image_name in enumerate(os.listdir(path)):
        if image_name.split('.')[-1] == 'png' or image_name.split('.')[-1] == 'jpg':
            image_path = os.path.join(path, image_name)
            xml_path = os.path.join(path, image_name.split('.')[0] + '.xml')
            img = cv_imread(image_path)
            boundingboxes = load_from_xml_file(xml_path)
            print(img.shape)
            seq = iaa.Sequential([
                iaa.Sharpen(alpha=alpha, lightness=lightness, name=None, deterministic=False, random_state=None)
            ])
            seq_det = seq.to_deterministic()  # 保持坐标和图像同步改变，而不是随机
            before_boundingBoxs = []
            for boundingbox in boundingboxes:
                # bndbox 坐标增强
                before_boundingBoxs.append(ia.BoundingBox(x1=boundingbox[0], y1=boundingbox[1]
                                                          , x2=boundingbox[2], y2=boundingbox[3]))

            bbs = ia.BoundingBoxesOnImage(before_boundingBoxs, shape=img.shape)
            bbs_aug = seq_det.augment_bounding_boxes([bbs])[0]
            image_aug = seq_det.augment_images([img])[0]
            after_boundingboxes = []

            for i in range(len(bbs.bounding_boxes)):
                after = bbs_aug.bounding_boxes[i]
                after_boundingboxes.append([int(after.x1), int(after.y1), int(after.x2), int(after.y2)])

            for i in range(len(bbs.bounding_boxes)):
                before = bbs.bounding_boxes[i]
                after = bbs_aug.bounding_boxes[i]
                print("BB %d: (%.4f, %.4f, %.4f, %.4f) -> (%.4f, %.4f, %.4f, %.4f)" % (
                    i,
                    before.x1, before.y1, before.x2, before.y2,
                    int(after.x1), int(after.y1), int(after.x2), int(after.y2))
                      )
            cv_imwrite(os.path.join(save_path, image_name), image_aug)
            shutil.copy(xml_path, os.path.join(save_path, image_name.split('.')[0] + '.xml'))
            change_from_xml_file(os.path.join(save_path, image_name.split('.')[0] + '.xml'), after_boundingboxes,
                                 image_aug.shape)


def contrastNormalization(path: str, save_path: str, alpha=1.0):
    '''
    改变图像的对比度
    :param path:要被缩放的路径
    :param save_path:缩放后保存的路径
    :param alpha: number or tuple of number or list of number or imgaug.parameters.StochasticParameter, optional
        Strength of the contrast normalization. Higher values than 1.0
        lead to higher contrast, lower values decrease the contrast.

            * If a number, then that value will be used for all images.
            * If a tuple ``(a, b)``, then a value will be sampled per image from
              the range ``a <= x <= b`` and be used as the alpha value.
            * If a list, then a random value will be sampled per image from
              that list.
            * If a StochasticParameter, then this parameter will be used to
              sample the alpha value per image.

    '''
    for idx, image_name in enumerate(os.listdir(path)):
        if image_name.split('.')[-1] == 'png' or image_name.split('.')[-1] == 'jpg':
            image_path = os.path.join(path, image_name)
            xml_path = os.path.join(path, image_name.split('.')[0] + '.xml')
            img = cv_imread(image_path)
            boundingboxes = load_from_xml_file(xml_path)
            print(img.shape)
            seq = iaa.Sequential([
                iaa.ContrastNormalization(alpha=alpha, per_channel=False, name=None, deterministic=False,
                                          random_state=None)
            ])
            seq_det = seq.to_deterministic()  # 保持坐标和图像同步改变，而不是随机
            before_boundingBoxs = []
            for boundingbox in boundingboxes:
                # bndbox 坐标增强
                before_boundingBoxs.append(ia.BoundingBox(x1=boundingbox[0], y1=boundingbox[1]
                                                          , x2=boundingbox[2], y2=boundingbox[3]))

            bbs = ia.BoundingBoxesOnImage(before_boundingBoxs, shape=img.shape)
            bbs_aug = seq_det.augment_bounding_boxes([bbs])[0]
            image_aug = seq_det.augment_images([img])[0]
            after_boundingboxes = []

            for i in range(len(bbs.bounding_boxes)):
                after = bbs_aug.bounding_boxes[i]
                after_boundingboxes.append([int(after.x1), int(after.y1), int(after.x2), int(after.y2)])

            for i in range(len(bbs.bounding_boxes)):
                before = bbs.bounding_boxes[i]
                after = bbs_aug.bounding_boxes[i]
                print("BB %d: (%.4f, %.4f, %.4f, %.4f) -> (%.4f, %.4f, %.4f, %.4f)" % (
                    i,
                    before.x1, before.y1, before.x2, before.y2,
                    int(after.x1), int(after.y1), int(after.x2), int(after.y2))
                      )
            cv_imwrite(os.path.join(save_path, image_name), image_aug)
            shutil.copy(xml_path, os.path.join(save_path, image_name.split('.')[0] + '.xml'))
            change_from_xml_file(os.path.join(save_path, image_name.split('.')[0] + '.xml'), after_boundingboxes,
                                 image_aug.shape)


if __name__ == "__main__":
    path = r'D:\360\4'
    save_path = r'D:\360\9'
    contrastNormalization(path, save_path, 2.0)
